package cas.ihep.spark.loop;

import cas.ihep.spark.io.ByteArrayOutputStreamExt;
import cas.ihep.spark.io.ByteBufferUtil;
import cas.ihep.spark.io.Packet;
import cas.ihep.spark.util.Counter;
import cas.ihep.spark.util.LogEnv;
import cas.ihep.spark.util.function.Function2;
import com.google.common.io.Closer;

import java.io.Closeable;
import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.util.*;

@SuppressWarnings("UnstableApiUsage")
public abstract class IterationNetServer implements Closeable{

    Counter<SocketAddress> numOfTasks;

    Closer closer;
    SocketAddress localAddress;//remoteAddress;
    Selector selector;

    private IterationNetServer(){
        closer=Closer.create();
    }

    public void close()throws IOException{
        closer.close();
    }

    public abstract void heartbeat()throws IOException;
    public abstract <T> T loop(Packet packet,int partitions,T zero,Function2<T,T,T> func,Class<T> cls)throws IOException;
    public <T> T iterate(Packet packet,int partitions,T zero,Function2<T,T,T> func)throws IOException{
        return zero;
    }
    public abstract void quit()throws IOException;

    public static IterationNetServer newServer(String type, SocketAddress local, Counter<SocketAddress> n)throws IOException{
        if("udp".equalsIgnoreCase(type)){
            return new UDPNet(local, n);
        }else if("tcp".equalsIgnoreCase(type)){
            return new TCPNet(local, n);
        }
        throw new IOException("Unsupported type of server: "+type);
    }

    private static class UDPNet extends IterationNetServer {
        DatagramChannel udp;
        Set<SocketAddress> replyHosts;
        HashMap<Integer,Object> objCache;
        UDPNet(SocketAddress local, Counter<SocketAddress> n)throws IOException{
            super();
            udp=closer.register(DatagramChannel.open());
            udp.bind(localAddress=local);
            //remoteAddress=remote;
            udp.configureBlocking(false);
            selector=closer.register(Selector.open());
            udp.register(selector, SelectionKey.OP_READ);
            numOfTasks=n;
            replyHosts=new HashSet<>();
            objCache=new HashMap<>();
            LogEnv.getInstance().debug("Creating a udp iteration net server");
        }

        public void heartbeat()throws IOException{
            LogEnv.getInstance().debug("Discovering all tasks");
            int desiredSize=numOfTasks.size();
            Collection<Map.Entry<SocketAddress,Integer>> numOfTasksValues=numOfTasks.values();
            HashMap<SocketAddress,Packet> packets=new HashMap<>();
            for(Map.Entry<SocketAddress,Integer> entry:numOfTasksValues){
                Packet packet=Packet.wrapEchoRequestPacket(entry.getValue());
                SocketAddress addr=entry.getKey();
                udp.send(packet.getBuffer(),addr);
                packets.put(addr,packet);
            }
            ByteBuffer buf=ByteBuffer.allocateDirect(128);
            try {
                while (replyHosts.size() < desiredSize) {
                    if (selector.select(1500) > 0) {
                        Iterator<SelectionKey> iter = selector.selectedKeys().iterator();
                        while (iter.hasNext()) {
                            SelectionKey key = iter.next();
                            iter.remove();
                            if (key.isReadable()) {
                                DatagramChannel dch = (DatagramChannel) key.channel();
                                SocketAddress sender = dch.receive(buf);
                                buf.flip();
                                Packet packet = Packet.wrap(buf.slice());
                                LogEnv.getInstance().debug("%s received", packet);
                                LogEnv.getInstance().debug("Desired packet<type: %d, reply: %d>", packet.type(), packet.echoReply());
                                if (packet.type() == Packet.EchoReply) {
                                    if (packet.echoReply() == 0) {
                                        replyHosts.add(sender);
                                        LogEnv.getInstance().debug("Get a echo ok reply from %s", sender);
                                    }
                                } else {
                                    LogEnv.getInstance().debug("Drop %s in heartbeat state", packet);
                                }
                                buf.clear();
                            }
                        }
                    } else {
                        for (Map.Entry<SocketAddress, Integer> entry : numOfTasksValues) {
                            SocketAddress addr = entry.getKey();
                            if (!replyHosts.contains(addr)) {
                                Packet packet = packets.get(addr);
                                assert (packet != null);
                                udp.send(packet.getBuffer(), addr);
                            }
                        }
                    }
                }
                LogEnv.getInstance().debug("All tasks have started");
            }finally {
                ByteBufferUtil.releaseBuffer(buf);
            }
        }

        @Override
        public <T> T loop(Packet packet, int partitions, T zeroValue, Function2<T,T,T> func,Class<T> klass) throws IOException {
            ByteBuffer sendBuf=packet.getBuffer();
            for(SocketAddress addr:replyHosts){
                udp.send(sendBuf,addr);
                sendBuf.rewind();
            }
            ByteBuffer buffer=ByteBuffer.allocateDirect(512);
            try {
                HashSet<SocketAddress> doesnotResponse = new HashSet<>(replyHosts);
                for (int j = 0; j < partitions; ) {
                    if (selector.select(500) > 0) {
                        Iterator<SelectionKey> iter = selector.selectedKeys().iterator();
                        while (iter.hasNext()) {
                            SelectionKey key = iter.next();
                            iter.remove();
                            if (key.isReadable()) {
                                DatagramChannel dch = (DatagramChannel) key.channel();
                                SocketAddress addr = dch.receive(buffer);
                                doesnotResponse.remove(addr);
                                buffer.flip();
                                Packet recvPacket = Packet.wrap(buffer.slice());
                                if (recvPacket.type() == Packet.DataReply) {
                                    j += recvPacket.taskCount();
                                    zeroValue = func.call(zeroValue, recvPacket.unwrapPacket(klass));
                                }
                                buffer.clear();
                            }
                        }
                    } else {
                        for (SocketAddress addr : doesnotResponse) {
                            udp.send(sendBuf, addr);
                            sendBuf.rewind();
                        }
                    }
                }
                return zeroValue;
            }finally {
                ByteBufferUtil.releaseBuffer(buffer);
            }
        }

        public <T> T iterate(Packet packet,int partitions,T zero,Function2<T,T,T> func)throws IOException{
            @SuppressWarnings("unchecked")
            Class<T> klass=(Class<T>)zero.getClass();
            int numOfIters=packet.numOfIterations();
            Object retObj;
            if((retObj=objCache.get(numOfIters))!=null){
                return klass.cast(retObj);
            }
            ByteBuffer sendBuf=packet.getBuffer();
            for(SocketAddress addr:replyHosts){
                udp.send(sendBuf,addr);
                sendBuf.rewind();
            }
            ByteBuffer buffer=ByteBuffer.allocateDirect(512);
            try {
                HashSet<SocketAddress> notResponse = new HashSet<>(replyHosts);
                for (int j = 0; j < partitions; ) {
                    if (selector.select(1000) > 0) {
                        Iterator<SelectionKey> iter = selector.selectedKeys().iterator();
                        while (iter.hasNext()) {
                            SelectionKey key = iter.next();
                            iter.remove();
                            if (key.isReadable()) {
                                DatagramChannel dch = (DatagramChannel) key.channel();
                                SocketAddress addr = dch.receive(buffer);
                                if (notResponse.contains(addr)) {
                                    buffer.flip();
                                    Packet recvPacket = Packet.wrap(buffer.slice());
                                    if (recvPacket.type() == Packet.DataReply) {
                                        if (recvPacket.numOfIterations() == numOfIters) {
                                            j += recvPacket.taskCount();
                                            zero = func.call(zero, recvPacket.unwrapPacket(klass));
                                            notResponse.remove(addr);
                                        } else {
                                            LogEnv.getInstance().debug("%s iteration number mismatch: actual %d, desired: %d",
                                                    recvPacket, recvPacket.numOfIterations(), numOfIters);
                                        }
                                    } else {
                                        LogEnv.getInstance().debug("Drop %s from %s in iterate state", recvPacket, addr);
                                    }
                                }
                                buffer.clear();
                            }
                        }
                    } else {
                        for (SocketAddress addr : notResponse) {
                            udp.send(sendBuf, addr);
                            sendBuf.rewind();
                        }
                    }
                }
                objCache.put(numOfIters, zero);
                return zero;
            }finally {
                ByteBufferUtil.releaseBuffer(buffer);
            }
        }

        @Override
        public void quit() throws IOException {
            Packet packet=Packet.QuitPacket();
            ByteBuffer buffer=packet.getBuffer();
            for(SocketAddress addr:replyHosts){
                udp.send(buffer,addr);
                buffer.rewind();
            }
        }
    }

    private static class TCPNet extends IterationNetServer {

        ServerSocketChannel tcp;
        ArrayList<SocketChannel> replyHosts;
        HashMap<SocketAddress, ByteArrayOutputStreamExt> cacheMap=new HashMap<>();
        HashMap<InetAddress, SocketAddress> addrCache=new HashMap<>();

        TCPNet(SocketAddress local, Counter<SocketAddress> n)throws IOException{
            super();
            localAddress=local;
            tcp=closer.register(ServerSocketChannel.open());
            tcp.bind(local);
            tcp.configureBlocking(false);
            selector=closer.register(Selector.open());
            tcp.register(selector,SelectionKey.OP_ACCEPT);
            numOfTasks=n;
            replyHosts=new ArrayList<>(5);
            LogEnv.getInstance().debug("Creating a tcp iteration net server");
        }

        private SocketAddress convertAddr(SocketAddress addr){
            InetSocketAddress iaddr=(InetSocketAddress)addr;
            InetAddress ia=iaddr.getAddress();
            SocketAddress ret=addrCache.get(ia);
            if(ret==null){
                ret=new InetSocketAddress(ia,8999);
                addrCache.put(ia,ret);
            }
            return ret;
        }

        @Override
        public void heartbeat() throws IOException {
            ArrayList<SocketChannel> unreplyHosts=new ArrayList<>(5);
            LogEnv.getInstance().debug("Discovering all tasks");
            int desiredSize=numOfTasks.size();
            //Collection<Map.Entry<SocketAddress,Integer>> numOfTasksValues=numOfTasks.values();
            ByteBuffer buffer=ByteBuffer.allocateDirect(128);
            ByteBuffer tmpBuf=null;
            try {
                while (replyHosts.size() < desiredSize) {
                    if (selector.select(500) > 0) {
                        Iterator<SelectionKey> iter = selector.selectedKeys().iterator();
                        while (iter.hasNext()) {
                            SelectionKey key = iter.next();
                            iter.remove();
                            if (key.isAcceptable()) {
                                ServerSocketChannel chn = (ServerSocketChannel) key.channel();
                                SocketChannel cli = closer.register(chn.accept());
                                cli.configureBlocking(false);
                                cli.register(selector, SelectionKey.OP_READ);
                                int cnt = numOfTasks.get(convertAddr(cli.getRemoteAddress()));
                                Packet packet = Packet.wrapEchoRequestPacket(cnt);
                                cli.write(packet.getBuffer());
                                LogEnv.getInstance().debug("One client from %s has connected", cli.getRemoteAddress());
                                unreplyHosts.add(cli);
                            } else if (key.isReadable()) {
                                SocketChannel chn = (SocketChannel) key.channel();
                                SocketAddress addr = chn.getRemoteAddress();
                                ByteArrayOutputStreamExt cache = cacheMap.get(addr);
                                if (cache == null) {
                                    cache = new ByteArrayOutputStreamExt(64);
                                    cacheMap.put(addr, cache);
                                }
                                chn.read(buffer);
                                buffer.flip();
                                int tmp = cache.size();
                                cache.write(buffer);
                                buffer.clear();
                                LogEnv.getInstance().debug("Reading %d bytes from %s", cache.size() - tmp, addr);
                                while (cache.size() > 0) {
                                    int bodyLength = cache.preadInt(0);
                                    int totalLength=bodyLength+4;
                                    if (cache.size() >= totalLength) {
                                        if(tmpBuf==null || tmpBuf.capacity()<totalLength){
                                            tmpBuf=ByteBuffer.allocate(totalLength);
                                        }
                                        tmpBuf.limit(totalLength);
                                        cache.readAndRemove(tmpBuf);
                                        tmpBuf.flip();
                                        Packet echo=Packet.wrap(tmpBuf);
                                        /*byte[] buf = new byte[bodyLength + 4];
                                        cache.readAndRemove(buf);
                                        Packet echo = Packet.wrap(buf);
                                        LogEnv.getInstance().debug("Reading a packet with %d(%d in total) bytes", bodyLength, bodyLength + 4);
                                        LogEnv.getInstance().debug("Desired packet<type: %d, reply: %d>", echo.type(), echo.echoReply());*/
                                        if (echo.type() == Packet.EchoReply && echo.echoReply() == 0) {
                                            replyHosts.add(chn);
                                            unreplyHosts.remove(chn);
                                            LogEnv.getInstance().debug("One packet<%s> received from %s", echo.toString(), addr);
                                        }
                                        tmpBuf.clear();
                                    } else {
                                        break;
                                    }
                                }
                            }
                        }
                    } else {
                        for (SocketChannel chn : unreplyHosts) {
                            int cnt = numOfTasks.get(convertAddr(chn.getRemoteAddress()));
                            Packet packet = Packet.wrapEchoRequestPacket(cnt);
                            chn.write(packet.getBuffer());
                        }
                    }
                }
                LogEnv.getInstance().debug("All tasks have started");
            }finally {
                ByteBufferUtil.releaseBuffer(buffer);
            }
        }

        public <T> T iterate(Packet packet,int partitions,T zero,Function2<T,T,T> func)throws IOException{
            @SuppressWarnings("unchecked")
            Class<T> cls=(Class<T>)zero.getClass();
            ByteBuffer sendBuf=packet.getBuffer();
            for(SocketChannel chn:replyHosts){
                chn.write(sendBuf);
                sendBuf.clear();
            }
            ByteBuffer buffer=ByteBuffer.allocateDirect(512);
            ByteBuffer tmpBuf=null;
            try {
                for (int j = 0; j < partitions; ) {
                    if (selector.select() > 0) {
                        Iterator<SelectionKey> iter = selector.selectedKeys().iterator();
                        while (iter.hasNext()) {
                            SelectionKey key = iter.next();
                            iter.remove();
                            if (key.isReadable()) {
                                SocketChannel chn = (SocketChannel) key.channel();
                                SocketAddress addr = chn.getRemoteAddress();
                                chn.read(buffer);
                                buffer.flip();
                                ByteArrayOutputStreamExt cache = cacheMap.get(addr);
                                if (cache == null) {
                                    cache = new ByteArrayOutputStreamExt(64);
                                    cacheMap.put(addr, cache);
                                }
                                cache.write(buffer);
                                buffer.clear();
                                while (cache.size() > 0) {
                                    int bodyLength = cache.preadInt(0);
                                    int totalLength=bodyLength+4;
                                    if (cache.size() >= totalLength) {
                                        if(tmpBuf==null || tmpBuf.capacity()<totalLength){
                                            tmpBuf=ByteBuffer.allocate(totalLength);
                                        }
                                        tmpBuf.limit(totalLength);
                                        cache.readAndRemove(tmpBuf);
                                        tmpBuf.flip();
                                        Packet recv=Packet.wrap(tmpBuf);
                                        /*byte[] buf = new byte[bodyLength + 4];
                                        cache.readAndRemove(buf);
                                        Packet recv = Packet.wrap(buf);*/
                                        if (recv.type() == Packet.DataReply) {
                                            j += recv.taskCount();
                                            zero = func.call(zero, recv.unwrapPacket(cls));
                                        }
                                        tmpBuf.clear();
                                    } else {
                                        break;
                                    }
                                }
                            }
                        }
                    }
                }
                return zero;
            }finally {
                ByteBufferUtil.releaseBuffer(buffer);
            }
        }

        @Override
        public <T> T loop(Packet packet, int partitions, T zero, Function2<T, T, T> func, Class<T> cls) throws IOException {
            ByteBuffer sendBuf=packet.getBuffer();
            for(SocketChannel chn:replyHosts){
                chn.write(sendBuf);
                sendBuf.clear();
            }
            ByteBuffer buffer=ByteBuffer.allocateDirect(512);
            ByteBuffer tmpBuf=null;
            try {
                for (int j = 0; j < partitions; ) {
                    if (selector.select() > 0) {
                        Iterator<SelectionKey> iter = selector.selectedKeys().iterator();
                        while (iter.hasNext()) {
                            SelectionKey key = iter.next();
                            iter.remove();
                            if (key.isReadable()) {
                                SocketChannel chn = (SocketChannel) key.channel();
                                SocketAddress addr = chn.getRemoteAddress();
                                chn.read(buffer);
                                buffer.flip();
                                ByteArrayOutputStreamExt cache = cacheMap.get(addr);
                                if (cache == null) {
                                    cache = new ByteArrayOutputStreamExt(64);
                                    cacheMap.put(addr, cache);
                                }
                                cache.write(buffer);
                                buffer.clear();
                                while (cache.size() > 0) {
                                    int bodyLength = cache.preadInt(0);
                                    int totalLength=bodyLength+4;
                                    if (cache.size() >= totalLength) {
                                        if(tmpBuf==null || tmpBuf.capacity()<totalLength){
                                            tmpBuf=ByteBuffer.allocate(totalLength);
                                        }
                                        tmpBuf.limit(totalLength);
                                        cache.readAndRemove(tmpBuf);
                                        tmpBuf.flip();
                                        Packet recv=Packet.wrap(tmpBuf);
                                        /*byte[] buf = new byte[bodyLength + 4];
                                        cache.readAndRemove(buf);
                                        Packet recv = Packet.wrap(buf);*/
                                        if (recv.type() == Packet.DataReply) {
                                            j += recv.taskCount();
                                            zero = func.call(zero, recv.unwrapPacket(cls));
                                        }
                                    } else {
                                        break;
                                    }
                                }
                            }
                        }
                    }
                }
                return zero;
            }finally {
                ByteBufferUtil.releaseBuffer(buffer);
            }
        }

        @Override
        public void quit() throws IOException {
            Packet packet=Packet.QuitPacket();
            ByteBuffer buf=packet.getBuffer();
            for(SocketChannel chn:replyHosts){
                chn.write(buf);
                buf.rewind();
            }
        }
    }
}
